import datetime

from bson import ObjectId

from models.base import (
    BaseModel,
)
from items.task import TaskStatus, TaskItem
from exceptions import ItemNotFoundException
from utils import logger
from database.redis_clients import RedisClients


class TaskModel(BaseModel):
    ret_item = TaskItem

    def __init__(self):
        super(TaskModel, self).__init__()
        self.sync_col.create_index([("create_time", 1)], expireAfterSeconds=86400 * 90, name='ttl_index')

    @RedisClients.lock("get_task_lock", 60)
    async def get(self, **kwargs) -> TaskItem | None:
        """
        获取一个任务
        """
        try:
            item = await self.async_find_one(
                {
                    "status": TaskStatus.ready_to_run,
                    "retry_times": {"$lt": 11},
                }
            )
        except ItemNotFoundException:
            return None
        else:
            print(f"获取任务成功{item.id}")
            await self.async_update_one(
                {"_id": item.id},
                {
                    "status": TaskStatus.running,
                    "update_time": datetime.datetime.now(),
                    "message": "正在运行",
                }
            )
            return item

    async def task_exits(
            self,
            platform,
            browser_id,
            exe_models,
            module_reference,
    ):
        """判断一个任务是否存在并未完成采集"""
        return await self.async_col.find_one({
            "platform": platform,
            "browser_id": browser_id,
            "exe_models": exe_models,
            "module_reference": module_reference,
            "status": {"$in": [TaskStatus.ready_to_run, TaskStatus.running]}
        })

    async def create(
            self,
            platform,
            browser_id,
            exe_models,
            module_reference,
            **kwargs
    ) -> TaskItem | None:
        """创建一个任务"""
        if await self.task_exits(
                platform,
                browser_id,
                exe_models,
                module_reference,
        ):
            logger.info("任务已存在且还没运行完成，无需重复创建")
            return None
        return await self.async_insert_one(self.ret_item(
            platform=platform,
            browser_id=browser_id,
            exe_models=exe_models,
            module_reference=module_reference,
            retry_times=0,
            status=TaskStatus.ready_to_run,
            message=""
        ))

    async def success(self, id: str, **kwargs) -> TaskItem:
        """
        任务执行成功
        """
        return await self.async_update_one(
            {"_id": ObjectId(id)},
            {
                "message": "采集成功",
                "status": TaskStatus.success,
                "update_time": datetime.datetime.now()
            }
        )

    async def failed(self, id, message, **kwargs) -> TaskItem:
        """任务执行失败"""
        return await self.async_update_one(
            {"_id": ObjectId(id)},
            {
                "message": message,
                "status": TaskStatus.failed,
                "update_time": datetime.datetime.now()
            }
        )

    async def retry(self, id, message, **kwargs) -> TaskItem:
        return await self.async_col.update_one(
            {"_id": ObjectId(id)},
            {
                "$set": {
                    "message": message,
                    "status": TaskStatus.ready_to_run,
                    "update_time": datetime.datetime.now()
                },
                "$inc": {"retry_times": 1}
            }
        )

    async def failed_model(self, id, model_name, message):
        ...


task_model = TaskModel()
